# -*- coding: utf-8 -*-
"""
Created on Mon Jan 28 14:41:40 2013

@author: daniel
"""

import scipy.optimize as opt
import numpy as np
import pp
class fmin:
    """
    Provides a wrapper for scipy optimizaiton frunction 'slsqp' that makes
    use of multi-core numerical gradient calculations.  Gradients for 
    objective and constraint functions are calculated using parallel forward
    differencing.  User provided objective and constarint functions must
    import all necessary libraries in order to run correctly since they will
    be launched in separate python environments.
    """
    def __init__(self,obj,X0,bounds,eqcons=[],ineqcons=[],args=(),dx=10**-7,nCPU=8,maxIter=100,iprint=1):
        self.X0=X0
        self.bounds=bounds
        self.setCPU(nCPU)
        self.DX=self.getDX(dx)
        self.Xmin=np.zeros(len(X0))
        self.fmin=0.0
        self.args=args
        self.maxIter=maxIter
        self.iprint=iprint
        self.objFcn=obj
        self.eqcons=eqcons
        self.ineqcons=ineqcons
        self.args=args
        self.setCPU(nCPU)
        self.xHistory=list()
        self.fHistory=list()
        if eqcons==None:
            self.deriv_eqcons=None
        else:
            self.deriv_eqcons=self._eqprime
        if ineqcons==None:
            self.deriv_ineqcons=None
        else:
            self.deriv_ineqcons=self._ineqprime
        if nCPU>1:self._derivFcn=self._gradient_Multi            
        else:self._derivFcn=self._gradient_Single
    def setCPU(self,N):
        self._nCPU=np.int(N) 
        if N>1:
            self._job_server=pp.Server(self._nCPU)        
        else:
            self._job_server=list()
    def getDX(self,dx):
        DX=np.zeros(len(self.X0))
        for i,bnd in enumerate(self.bounds):
            DX[i]=(bnd[1]-bnd[0])*dx
        return DX
    def solve(self):
        result=opt.fmin_slsqp(self.objFcn,self.X0,eqcons=self.eqcons,ieqcons=self.ineqcons,bounds=self.bounds,fprime=self._fprime,fprime_eqcons=self.deriv_eqcons,fprime_ieqcons=self.deriv_ineqcons,args=self.args,iter=self.maxIter,iprint=self.iprint)
        self.Xmin=result
        self.fmin=self.objFcn(self.Xmin,*self.args)
    def _gradient_Multi(self,X,dx,fcn,args):
        X=np.array(X)
        N=len(X)
        fprime=np.zeros(N)
        jobs=[]
        fX=fcn(X,*args)
        self.xHistory.append(X)
        self.fHistory.append(fX)
        for i in range(N):
            XdX=X.copy()
            XdX[i]=XdX[i]+dx[i]
            inp=self._getInp(XdX,args)
            jobs.append(self._job_server.submit(fcn,inp))
        for i,job in enumerate(jobs):
            completedJob=job()
            fXdX=completedJob            
            fprime[i]=(fXdX-fX)/dx[i]
        return fprime
    def _gradient_Single(self,X,dx,fcn,args):
        X=np.array(X)
        N=len(X)
        fprime=np.zeros(N)
        fX=fcn(X,*args)
        self.xHistory.append(X)
        self.fHistory.append(fX)
        for i in range(N):
            XdX=X.copy()
            XdX[i]=XdX[i]+dx[i]
            fXdX=fcn(XdX,*args)
            fprime[i]=(fXdX-fX)/dx[i]
        return fprime
    def _getInp(self,X,args):
        inp=list()
        inp.append(X)
        for arg in args:inp.append(arg)
        return tuple(inp)
    def _fprime(self,*args):
        X=args[0]
        return self._derivFcn(X,self.DX,self.objFcn,self.args)
    def _eqprime(self,X):
        gp=np.zeros([len(self.eqcons),len(X)])
        for i,con in enumerate(self.eqcons):
            gp[i,:]=self._derivFcn(X,self.DX,con,self.args)
    def _ineqprime(self,X):
        gp=np.zeros([len(self.ineqcons),len(X)])
        for i,con in enumerate(self.ineqcons):
            gp[i,:]=self._derivFcn(X,self.DX,con,self.args)
    def _denorm(self,Xn):
        return (self.ub-self.lb)*Xn+self.lb
def fprimeTest():
    obj=rosenObj
    ineqcon=list()
    ineqcon.append(rosenCon)
    arg=()#tuple([1,2,3])
    x0=[5.1,-5.2]
    bounds=[[-10,10],[-10,10]]
    import time
    time1=time.time()
    fm=fmin(obj,x0,bounds,args=arg,ineqcons=ineqcon,nCPU=8)
    fm.solve()
    time2=time.time()
    fm=fmin(obj,x0,bounds,args=arg,ineqcons=ineqcon,nCPU=1)
    fm.solve()
    time3=time.time()
    print time1
    print time2-time1
    print time3-time2
    #print fm._gradient(x0,0.001,obj,args=arg)
    
def rosenObj(x):
    for i in range(100000):
        f=100.*(x[1]-x[0]**2)**2+(1-x[0])**2
    return f
def rosenCon(x):
    g=x[0]-2.*x[1]
    return g
 
if __name__=="__main__":   
    fprimeTest()